package org.dromara.common.websocket.utils

import cn.hutool.core.collection.CollUtil
import io.github.oshai.kotlinlogging.KotlinLogging
import org.dromara.common.redis.utils.RedisUtils.publish
import org.dromara.common.redis.utils.RedisUtils.subscribe
import org.dromara.common.websocket.constant.WebSocketConstants.WEB_SOCKET_TOPIC
import org.dromara.common.websocket.dto.WebSocketMessageDto
import org.dromara.common.websocket.holder.WebSocketSessionHolder
import org.springframework.web.socket.PongMessage
import org.springframework.web.socket.TextMessage
import org.springframework.web.socket.WebSocketMessage
import org.springframework.web.socket.WebSocketSession
import java.io.IOException
import java.util.function.Consumer

/**
 * 工具类
 *
 * @author zendwang
 * @author LikeYouDo
 * @date 2025/1/7 10:20
 */
object WebSocketUtils {
    private val log = KotlinLogging.logger {}

    /**
     * 向指定的WebSocket会话发送消息
     *
     * @param sessionKey 要发送消息的用户id
     * @param message    要发送的消息内容
     */
    @JvmStatic
    fun sendMessage(sessionKey: Long, message: String) {
        val session = WebSocketSessionHolder.getSessions(sessionKey)
        session?.let { sendMessage(it, message) }
    }

    /**
     * 订阅WebSocket消息主题，并提供一个消费者函数来处理接收到的消息
     *
     * @param consumer 处理WebSocket消息的消费者函数
     */
    @JvmStatic
    fun subscribeMessage(consumer: Consumer<WebSocketMessageDto>) {
        subscribe(WEB_SOCKET_TOPIC, WebSocketMessageDto::class.java, consumer)
    }

    /**
     * 发布WebSocket订阅消息
     *
     * @param webSocketMessage 要发布的WebSocket消息对象
     */
    @JvmStatic
    fun publishMessage(webSocketMessage: WebSocketMessageDto) {
        val unsentSessionKeys: MutableList<Long> = ArrayList()
        // 当前服务内session,直接发送消息
        for (sessionKey in webSocketMessage.sessionKeys!!) {
            if (WebSocketSessionHolder.existSession(sessionKey)) {
                sendMessage(sessionKey, webSocketMessage.message!!)
                continue
            }
            unsentSessionKeys.add(sessionKey)
        }
        // 不在当前服务内session,发布订阅消息
        if (CollUtil.isNotEmpty(unsentSessionKeys)) {
            val broadcastMessage = WebSocketMessageDto().apply {
                this.message = webSocketMessage.message
                sessionKeys = unsentSessionKeys
            }
            publish(WEB_SOCKET_TOPIC, broadcastMessage) {
                log.info {
                    " WebSocket发送主题订阅消息topic:$WEB_SOCKET_TOPIC session keys:$unsentSessionKeys message:$webSocketMessage.message"
                }
            }
        }
    }

    /**
     * 向所有的WebSocket会话发布订阅的消息(群发)
     *
     * @param message 要发布的消息内容
     */
    @JvmStatic
    fun publishAll(message: String) {
        val broadcastMessage = WebSocketMessageDto().apply {
            this.message = message
        }
        publish(WEB_SOCKET_TOPIC, broadcastMessage) {
            log.info {
                "WebSocket发送主题订阅消息topic:$WEB_SOCKET_TOPIC message:$message"
            }
        }
    }

    /**
     * 向指定的WebSocket会话发送Pong消息
     *
     * @param session 要发送Pong消息的WebSocket会话
     */
    @JvmStatic
    fun sendPongMessage(session: WebSocketSession) {
        sendMessage(session, PongMessage())
    }

    /**
     * 向指定的WebSocket会话发送文本消息
     *
     * @param session WebSocket会话
     * @param message 要发送的文本消息内容
     */
    @JvmStatic
    fun sendMessage(session: WebSocketSession, message: String) {
        sendMessage(session, TextMessage(message))
    }

    /**
     * 向指定的WebSocket会话发送WebSocket消息对象
     *
     * @param session WebSocket会话
     * @param message 要发送的WebSocket消息对象
     */
    @JvmStatic
    @Synchronized
    private fun sendMessage(session: WebSocketSession, message: WebSocketMessage<*>) {
        if (!session.isOpen) {
            log.warn { "[send] session会话已经关闭" }
        } else {
            try {
                session.sendMessage(message)
            } catch (e: IOException) {
                log.error(e) { "[send] session($session) 发送消息($message) 异常" }
            }
        }
    }
}
